# -*- coding: utf-8 -*-#

# -------------------------------------------------------------------------------
# Name:         momentum_stratege.py
# Description:  动量策略
# Author:       ylf
# Date:         9/23/21
# Detail:       筛选沪深300中符合条件的股票，计算动量因子和交易信号，根据等权重计算收益率
# -------------------------------------------------------------------------------
import data.stock as st
import pandas as pd
import strategy.base as base


def get_data(
        stock_id,
        start_date,
        end_date,
        length=20,
        type='price',
        columns=['date', 'close']):
    """
    获取到前10支股票的
    :return:
    """
    stocks = st.get_index_stocks(stock_id)
    return_df = pd.DataFrame()
    # 获取需要的数据长度，默认为全部
    if length == 'all':
        use_lenth = len(stocks)
    else:
        use_lenth = length
    for stock in stocks[:use_lenth]:
        st_data = st.get_csv_price(
            stock,
            start_date,
            end_date,
            type=type,
            columns=columns)
        st_data.index = pd.to_datetime(st_data.index)
        st_data = st_data.resample('M').last()
        # 生成所需要的数据并进行返回
        st_data['month_earning'] = st_data['close'] / \
            st_data['close'].shift() - 1
        st_data.drop(columns=['close'], axis=1, inplace=True)
        st_data.rename(columns={'month_earning': stock}, inplace=True)
        return_df = pd.concat([return_df, st_data], axis=1)
    return return_df


def get_signals(data, top_n=2):
    """
    极值信号寻找工具【反向寻找可对data*-1,然后进行相加操作】
    :param data: 传入的数据
    :param top_n: 产生的信号个数
    :return: 0-1信号表
    """
    # 创建一个和data相同的dataframe，包含行名和列名，数据值为空
    signals = pd.DataFrame(data=None, index=data.index, columns=data.columns, dtype=object)
    # # 循环数据
    for index, row in data.iterrows():
        signals.loc[index] = row.isin(row.nlargest(top_n)).astype(int)
    return signals


def momentum_stratege(st_data, top_n=2):
    """
    动量策略
    :return:
    """
    # 1 生成交易信号
    # 将st_data的Nan进行处理
    st_data.fillna(0, inplace=True)
    reverse_st_data = -1 * st_data
    buy_signals = get_signals(st_data, top_n=2)
    sell_signals = get_signals(reverse_st_data, top_n=2)
    final_signals = buy_signals - sell_signals
    print(final_signals)
    # 2 计算投资组合收益率
    returns = base.caculate_combination_rate(st_data, final_signals, 2 * top_n)
    print(returns)

    # 3 评估策略效果
    # 4 数据预览
if __name__ == '__main__':
    """
    调用主策略
    """
    # 首先获取数据

    return_df = get_data(
        '000300.XSHG',
        '2020-01-01',
        '2021-09-01',
        length=5,
        columns=[
            'date',
            'close'])
    print(return_df)
    momentum_stratege(return_df)
